import numpy as np
from math import sqrt
from collections import Counter

class KNNClassifier:

    def __init__(self, k):
        assert k >= 1, "k must be valid"
        self.k = k
        self.xTrain = None
        self.yTrain = None

    def fit(self, xTrain, yTrain):
        assert xTrain.shape[0] == yTrain.shape[0], \
            "the size of xTrain must be equal to the size of yTrain"
        assert self.k <= xTrain.shape[0], \
            "the size of xTrain must be at least k"

        self.xTrain = xTrain
        self.yTrain = yTrain
        return self

    def predict(self, xPredict):
        assert self.xTrain is not None and self.yTrain is not None, \
            "must fit before predict"
        assert xPredict.shape[1] == self.xTrain.shape[1], \
            "the feature number of xPredict must be equal to xTrain"

        yPredict = [self._predict(x) for x in xPredict]
        return np.array(yPredict)

    def _predict(self, x):
        # assert x.shape[0] == self.xTrain.shape[0],\
        #     "the featur number of x must be equal to xTrain"

        distances = [sqrt(np.sum((xTrain - x) ** 2)) for xTrain in self.xTrain]

        nearest = np.argsort(distances)
        topY = [self.yTrain[i] for i in nearest[:self.k]]

        votes = Counter(topY)

        return votes.most_common(1)[0][0]

    def __repr__(self):
        return "KNN=(k=%d)" % self.k

    def trainTestSplit(x, y, testRatio=0.2, seed=None):
        assert x.shape[0] == y.shape[0], \
            'the size of x must b equal to the size of y'

        assert 0.0 <= testRatio <= 1.0, \
            'test ratio must be valid'

        if seed:
            np.random.seed(seed)

        shuffleIndexes = np.random.permutation(len(x))

        testSize = int(len(x) * testRatio)
        testIndexes = shuffleIndexes[:testSize]

        trainIndexes = shuffleIndexes[testSize:]

        xTrain = x[trainIndexes]
        yTrain = y[trainIndexes]

        xTest = x[testIndexes]
        yTest = y[testIndexes]

        return xTrain, xTest, yTrain, yTest

    def accurayScore(self, yTrue, yPredict):
        assert yTrue.shape[0] == yPredict.shape[0], \
            'the size of yTrue must be equal to the size of yPredict'

        return sum(yTrue == yPredict) / len(yTrue)

    def score(self, xTest, yTest):
        yPredict = self.predict(xTest)
        return self.accurayScore(yTest, yPredict)